/*
 * Copyright 2021 Rikai authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.spark.sql.rikai

import org.apache.spark.sql.{DataFrame, SparkSession}
import org.scalatest.{BeforeAndAfterAll, BeforeAndAfterEach, Suite}
import org.slf4j.LoggerFactory
import ch.qos.logback.classic.{Level, Logger}
import scala.util.Properties

trait SparkTestSession extends BeforeAndAfterEach with BeforeAndAfterAll {

  this: Suite =>

  private val level = Level.valueOf {
    Properties.envOrElse("LOG_LEVEL", Level.ERROR.levelStr)
  }
  LoggerFactory
    .getLogger("root")
    .asInstanceOf[Logger]
    .setLevel(level)
  LoggerFactory
    .getLogger("org.apache.spark.sql.rikai")
    .asInstanceOf[Logger]
    .setLevel(Level.DEBUG)

  lazy val spark: SparkSession = SparkSession.builder
    .config(
      "spark.sql.extensions",
      "org.apache.spark.sql.rikai.LigaImageExtensions"
    )
    .config("spark.port.maxRetries", 128)
    .master("local[*]")
    .getOrCreate

  spark.sparkContext.setLogLevel("WARN")

  def assertEqual(actual: DataFrame, expected: DataFrame): Unit = {
    assert(actual.count() == expected.count())
    assert(actual.exceptAll(expected).isEmpty)
  }

  override protected def afterAll(): Unit = {
    spark.close()
  }
}
